// gcc -O3 -I. cross_powerspc2.c util.c -lfftw3f -lm -o cross_powerspc2

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <complex.h>
#include <fftw3.h>
#include <limits.h>

float hypot3(float x, float y, float z);
float dot(float x[3], float y[3]);
float **matrix(long nrow, long ncol);
float ***tensor3(int n1, int n2, int n3);
float complex ***ctensor3(int n1, int n2, int n3);
int get_option(int argc, char *argv[], char format[],...);

// find the grid cells covered by particle at x & compute the CIC weights
void CIC_weights(float x[], int ng, float *xa, int ind[3][2], float w[3][2])
{
  int i;
  float dx = xa[1] - xa[0];

  for (i=0; i<3; i++) {
    // first make sure that positions fall in box volume
    while (x[i] < xa[0]-dx/2) x[i] += ng*dx;
    while (x[i] > xa[ng-1]+dx/2) x[i] -= ng*dx;

    // now figure out nearest grid cells
    ind[i][0] = (int)floor((x[i]-xa[0])/dx);
    if (ind[i][0]<0 || ind[i][0]>=ng-1) {
      ind[i][0]=ng-1; ind[i][1]=0;
      if (x[i]<xa[ind[i][0]]) 
	w[i][1] = 1 - (w[i][0] = (xa[ind[i][1]]-x[i])/dx); 
      else w[i][0] = 1 - (w[i][1]=(x[i]-xa[ind[i][0]])/dx);
    } else {
      ind[i][1]=ind[i][0]+1; 
      w[i][0] = 1 - (w[i][1]=(x[i]-xa[ind[i][0]])/dx);
    }
  }

  return;
}

float sinc(float x) {return (fabs(x)<1e-5 ? 1.0 : sin(x)/x);}

float sqr(float x) {return x*x;}
    
float FT_CIC(float k[3], float dx)
{return sqr(sinc(k[0]*dx/2)*sinc(k[1]*dx/2)*sinc(k[2]*dx/2));}

FILE *fopen_check(char *filename, char *mode)
{
  FILE *fp = fopen(filename, mode);
  if (fp == NULL) {
    fprintf(stderr,"Error opening file %s with mode %s\n",filename,mode);
    exit(1);
  }
  return fp;
}

void get_crossPk(int n1, float xa[], float ***g1, float ***g2, 
		 int nk, char fileroot[], int nh)
{
  int i,j,l, bin;
  long ng=n1, ntot;
  fftwf_complex ***Fg1=NULL, ***Fg2=NULL;
  fftwf_plan plan;
  float dx, xsize, knorm, k[3], dk, kmax, dP; 
  int countk[nk];
  float ka[nk], Pk[nk], sig[nk];
  float Pm[nk], Pg[nk], Pc[nk];
  char string[500], st[500];
  FILE *fp;

  dx = xa[1]-xa[0];
  xsize = dx*ng;
  dk = 2.0*M_PI / xsize;
  kmax = M_PI / dx;
  ntot = ng*ng*ng;

//  sprintf(st,"In getPk ng=%ld, dx=%g, xsize=%g, dk=%g, kmax=%g, ntot=%ld\n",ng,dx,xsize,dk,kmax,ntot);
//  printf("%s",st);
  fprintf(stdout,"In getPk ng=%ld, dx=%g, xsize=%g, dk=%g, kmax=%g, ntot=%ld\n",ng,dx,xsize,dk,kmax,ntot);
// fputs(st,stderr); fflush(stderr);

  Fg1 = ctensor3(ng,ng,ng);
  plan = fftwf_plan_dft_r2c_3d(ng,ng,ng,**g1,**Fg1,FFTW_ESTIMATE);  
  fftwf_execute(plan);
  fftwf_destroy_plan(plan);
  free(**g1); free(*g1); free(g1); g1=NULL;

  if (g2) {
    Fg2 = ctensor3(ng,ng,ng);
    plan = fftwf_plan_dft_r2c_3d(ng,ng,ng,**g2,**Fg2,FFTW_ESTIMATE);  
    fftwf_execute(plan);
    fftwf_destroy_plan(plan);
    free(**g2); free(*g2); free(g2); g2=NULL;
  }

  // first matter auto-spectrum
  sprintf(string, "%s.mpk", fileroot);
  fp = fopen_check(string, "w");
  for (i=0; i<nk; i++) {
    Pk[i] = sig[i] = countk[i] = 0;
    ka[i] = i * kmax / nk;
  }
  for (i=0; i<ng; i++) for (j=0; j<ng; j++) for (l=0; l<ng/2 + 1; l++) {
    k[0] = dk * (i < (ng/2 + 1) ? i : i-ng);
    k[1] = dk * (j < (ng/2 + 1) ? j : j-ng);
    k[2] = dk*l;
    knorm = sqrt(dot(k,k));
    bin = (int)rint(knorm*nk/kmax);
    if (bin>=0 && bin<nk) {
      ++(countk[bin]);
      dP = sqr(cabs(Fg1[i][j][l])/FT_CIC(k,dx));
      Pk[bin] += dP;
      sig[bin] += sqr(dP);
    }
  }
  for (i=1; i<nk; i++) if (countk[i]) {
    Pk[i] *= dx*dx*dx/ntot/countk[i];
    sig[i] *= sqr(dx*dx*dx/ntot)/countk[i];
    sig[i] = sqrt((sig[i]-sqr(Pk[i]))/(countk[i]-1));
    fprintf(fp, "%g\t%g\t%g\n", ka[i]/1200., Pk[i]*1200.*1200.*1200., sig[i]*1200.*1200.*1200.);
    Pm[i] = Pk[i];
  }
  fclose(fp);

  if (Fg2) {
  // next halo auto-spectrum
    sprintf(string, "%s.gpk", fileroot);
    fp = fopen_check(string, "w");
    for (i=0; i<nk; i++) Pk[i] = sig[i] = countk[i] = 0;
    for (i=0; i<ng; i++) for (j=0; j<ng; j++) for (l=0; l<ng/2 + 1; l++) {
      k[0] = dk * (i < (ng/2 + 1) ? i : i-ng);
      k[1] = dk * (j < (ng/2 + 1) ? j : j-ng);
      k[2] = dk*l;
      knorm = sqrt(dot(k,k));
      bin = (int)rint(knorm*nk/kmax);
      if (bin>=0 && bin<nk) {
	++(countk[bin]);
	dP = sqr(cabs(Fg2[i][j][l])/FT_CIC(k,dx));
	Pk[bin] += dP;
	sig[bin] += sqr(dP);
      }
    }
    for (i=1; i<nk; i++) if (countk[i]) {
      Pk[i] *= dx*dx*dx/ntot/countk[i];
      sig[i] *= sqr(dx*dx*dx/ntot)/countk[i];
      sig[i] = 1200.*1200.*1200.*sqrt((sig[i]-sqr(Pk[i]-1.0/nh))/(countk[i]-1));
      fprintf(fp, "%g\t%g\t%g\n", ka[i]/1200., (Pk[i]-1.0/nh)*1200.*1200.*1200., sig[i]);
      Pg[i] = (Pk[i]-1.0/nh);
    }
    fclose(fp);
    // last the cross-spectrum
    sprintf(string, "%s.cpk", fileroot);
    fp = fopen_check(string, "w");
    for (i=0; i<nk; i++) Pk[i] = sig[i] = countk[i] = 0;
    for (i=0; i<ng; i++) for (j=0; j<ng; j++) for (l=0; l<ng/2 + 1; l++) {
      k[0] = dk * (i < (ng/2 + 1) ? i : i-ng);
      k[1] = dk * (j < (ng/2 + 1) ? j : j-ng);
      k[2] = dk*l;
      knorm = sqrt(dot(k,k));
      bin = (int)rint(knorm*nk/kmax);
      if (bin>=0 && bin<nk) {
	++(countk[bin]);
	dP = creal(Fg1[i][j][l]*conj(Fg2[i][j][l]))/sqr(FT_CIC(k,dx));
	Pk[bin] += dP;
	sig[bin] += sqr(dP);
      }
    }
    for (i=1; i<nk; i++) if (countk[i]) {
      //printf("ka: %g, Pk: %g, dx: %g, ntot: %d, count: %d\n",ka[i], Pk[i], dx, ntot, countk[i]);      
      Pk[i] *= dx*dx*dx/ntot/countk[i];
      sig[i] *= sqr(dx*dx*dx/ntot)/countk[i];
      sig[i] = 1200.*1200.*1200.*sqrt((sig[i]-sqr(Pk[i]))/(countk[i]-1));
      fprintf(fp, "%g\t%g\t%g\n", ka[i]/1200., Pk[i]*1200.*1200.*1200., sig[i]);
      Pc[i] = Pk[i];
    }
    fclose(fp);
    free(**Fg2); free(*Fg2); free(Fg2);

    sprintf(string, "%s.rpk", fileroot);
    fp = fopen_check(string, "w");
    for (i=1; i<nk; i++) if (countk[i]) 
      fprintf(fp, "%g\t%g\t%g\n", ka[i]/1200., Pc[i]/sqrt(Pm[i]*Pg[i]), 
	      sig[i]/sqrt(Pm[i]*Pg[i]));
    fclose(fp);
  }
  free(**Fg1); free(*Fg1); free(Fg1);

  return;
}

int main(int argc, char *argv[])
{
  int nh=-1, ng=512, ng3;
  float ***delall=NULL, ***delgrp=NULL, *xa, x[3], d=0, d0=0;
  float w[3][2];
  int p, i, j, k, ind[3][2]; 
  int nk, RENORM=0;
  FILE *fp;
  char string[500], outfile[100]="", delfile[100], catalogfile[100];

  if (!get_option(argc, argv, "-d %s", delfile)) {
    fputs("Must specify delta file with -d\n",stderr);
    exit(EXIT_FAILURE);
  }
  if (!get_option(argc, argv, "-c %s", catalogfile)) {
    fputs("Must specify catalog file with -c\n",stderr);
    exit(EXIT_FAILURE);
  }
/*** 
 if (!get_option(argc, argv, "-o %s", outfile)) {
    fputs("Must specify output file with -o\n",stderr);
    exit(EXIT_FAILURE);
  }
***/
  RENORM = get_option(argc, argv, "-renorm");

  fp = fopen_check(delfile, "rb");
  fseek(fp, 0, SEEK_END);
  ng3 = ftell(fp)/4;
  ng = rint(cbrt(ng3));
  if (ng*ng*ng != ng3) {
    fprintf(stderr, "Error: appear to be %d floats in %s, not a cube?\n",
	    ng3,delfile); exit(EXIT_FAILURE);
  }
  xa = calloc(ng,4);
  for (i=0; i<ng; i++) xa[i] = (i+0.5)/ng;
  // zero the density grids
  delall=tensor3(ng,ng,ng);
  rewind(fp);
  fread(**delall, 4, ng3, fp);
  fclose(fp);

  fprintf(stderr,"found ng3=%d ng=%d\n",ng3, ng); 

  delgrp=tensor3(ng,ng,ng);
  fp = fopen_check(catalogfile, "r");
  for (d=0, nh=-1; !feof(fp); nh++) if (fgets(string, 499, fp)) {
    sscanf(string, "%f", &d0);
    d += d0;
  }
  printf("# of halos: %d\n", nh);
  rewind(fp);
  d0 = d/ng3;
  if (RENORM) for (p=0; p<ng3; p++) (**delgrp)[p] = -1;
  else for (p=0; p<ng3; p++) (**delgrp)[p] = -d0;

  while (!feof(fp)) if (fgets(string, 499, fp)) {
    sscanf(string, "%f %f %f %f", &d, x, x+1, x+2);
    if (RENORM) d /= d0;
    CIC_weights(x,ng,xa,ind,w);
    // now add densities to cells with appropriate weights
    for (i=0; i<2; i++) for (j=0; j<2; j++) for (k=0; k<2; k++) 
      delgrp[ind[0][i]][ind[1][j]][ind[2][k]] += d*w[0][i]*w[1][j]*w[2][k];
  }
  fclose(fp);
  nk = ng;

 // instead of calling the get_crossPk, make 1D density maps
  FILE *dm;
  FILE *gm;
  dm=fopen("dm.txt", "w");
  gm=fopen("gm.txt", "w");
 for (i=0; i<ng; i++) {
      fprintf(dm, "%i %g\n", i, delall[i][0][0]);
      fprintf(gm, "%i %g\n", i, delgrp[i][0][0]);
}
fclose(dm);
fclose(gm);
dm=fopen("dm2d.txt", "w");
gm=fopen("gm2d.txt", "w");

// write a 2d density map
for (i=0; i<ng; i++) {
  for (j=0; j<ng; j++) {
    fprintf(dm, "%i %i %g\n", i, j, delall[i][j][0]);
    fprintf(gm, "%i %i %g\n", i, j, delgrp[i][j][0]);
  }
}


return 0;
  fputs("calling get_Pk\n",stderr); fflush(stderr);
  get_crossPk(ng, xa, delall, delgrp, nk, outfile, nh);
  fputs("finished get_Pk\n",stderr); fflush(stderr);

  return 0;
}
